#include <stdio.h>
#include <algorithm>
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>

#define THREADS_PER_BLOCK 2
#define MAX_BLOCKS 2
#define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0))
#define CUDA_2D_KERNEL_LOOP(idx, numx, idy, numy)     \
    for(size_t idx = blockIdx.x * blockDim.x + threadIdx.x; \
        idx < (numx); idx += blockDim.x*gridDim.x)    \
        for(size_t idy = blockIdx.y * blockDim.y + threadIdx.y; \
            idy < (numy); idy += blockDim.y*gridDim.y)    \


inline int get_blocks(int total_threads , int threads_one_block){
    int blocks = (total_threads + threads_one_block -1) /threads_one_block;
    blocks = std::min(blocks, MAX_BLOCKS);
    return blocks;
}

//#define DEBUG
const int THREADS_PER_BLOCK_NMS = sizeof(unsigned long long) * 8;
__device__ const float EPS = 1e-8;

__device__ float iou_core(const float* box1 ,const float* box2){
    float box1_x0 = *(box1 + 0);
    float box1_y0 = *(box1 + 1);
    float box1_x1 = *(box1 + 2);
    float box1_y1 = *(box1 + 3);
    float box2_x0 = *(box2 + 0);
    float box2_y0 = *(box2 + 1);
    float box2_x1 = *(box2 + 2);
    float box2_y1 = *(box2 + 3);
    if(!(box1_x0 < box1_x1 && box1_y0 < box1_y1 && box2_x0 < box2_x1 && box2_y0 < box2_y1)){
        return 9999;
    }

    float inter_x0 = std::max(box1_x0, box2_x0);
    float inter_x1 = std::min(box1_x1, box2_x1);
    float inter_y0 = std::max(box1_y0, box2_y0);
    float inter_y1 = std::min(box1_y1, box2_y1);
    float inter_area = (inter_x1 - inter_x0)*(inter_y1-inter_y0);
    inter_area = std::max(inter_area, 0.0f);

    float box1_area = (box1_x1 - box1_x0)*(box1_y1-box1_y0);
    float box2_area = (box2_x1 - box2_x0)*(box2_y1-box2_y0);
    float iou = inter_area / (box1_area + box2_area - inter_area);
//     printf("iou =%f\n",iou);
    return iou;
}

__global__ void iou_gpu_kernel(const int box1_num,
const float* box1_ptr,
const int box2_num,
const float* box2_ptr,
float* result_ptr){
    CUDA_2D_KERNEL_LOOP(box1_idx, box1_num, box2_idx, box2_num)
    {
//         const int box1_idx = blockIdx.x * THREADS_PER_BLOCK + threadIdx.x;
//         const int box2_idx = blockIdx.y * THREADS_PER_BLOCK + threadIdx.y;
        printf("gpu: box1_idx = %d, box2_idy= %d\n",box1_idx,box2_idx);
        if(box1_idx>=box1_num || box2_idx>=box2_num){
            return;
        }
        printf("gpu: box1_idx = %d, box2_idy= %d, result_id= %d\n",box1_idx,box2_idx,box1_idx * box2_num + box2_idx);
        const float* box1 = box1_ptr + box1_idx * 4;
        const float* box2 = box2_ptr + box2_idx * 4;
        float iou = iou_core(box1, box2);
        *(result_ptr + box1_idx * box2_num + box2_idx) = iou;
    }
}

void iou_gpu_launch(const int box1_num,
const float* box1_ptr,
const int box2_num,
const float* box2_ptr,
float* result_ptr){
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    dim3 blocks(get_blocks(box1_num, THREADS_PER_BLOCK),get_blocks(box2_num, THREADS_PER_BLOCK));//每个grid的blocks
    dim3 threads(THREADS_PER_BLOCK,THREADS_PER_BLOCK);//每个block里面的thread
    printf("blocks=(%d %d), threads=(%d %d)\n",
        get_blocks(box1_num, THREADS_PER_BLOCK),get_blocks(box2_num, THREADS_PER_BLOCK),
        THREADS_PER_BLOCK,THREADS_PER_BLOCK);
    iou_gpu_kernel<<<blocks,threads,0,stream>>>(box1_num,box1_ptr,box2_num,box2_ptr,result_ptr);
    cudaDeviceSynchronize();// waiting for gpu work
    printf("gpu done\n");
}
